package com.ruibo;

import org.pytorch.IValue;
import org.pytorch.Module;
import org.pytorch.Tensor;

/**
 * PyTorchModelExample
 *
 * @author sunfeng
 * @date 2025/2/22 11:40
 */
public class PyTorchModelExample {
    public static void main(String[] args) {
        try {
            // 加载PyTorch模型
            Module module = Module.load("model.pt");

            // 创建输入Tensor
            float[] inputData = {1.0f, 2.0f, 3.0f, 4.0f};
            Tensor inputTensor = Tensor.fromBlob(inputData, new long[]{1, 4});

            // 运行模型
            IValue output = module.forward(IValue.from(inputTensor));

            // 获取输出Tensor
            Tensor outputTensor = output.toTensor();

            // 打印输出
            float[] outputData = outputTensor.getDataAsFloatArray();
            for (float value : outputData) {
                System.out.println(value);
            }

        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}
